from flask import Flask, request, jsonify, render_template
from knowledge_base import save_uploaded_file, list_uploaded_files, load_selected_files
import json
import requests
import re
import logging
import os

app = Flask(__name__)

DEEPSEEK_API_KEY = "请填写密钥"
DEEPSEEK_API_URL = "https://api.deepseek.com/v1/chat/completions"

# 设置日志文件
logging.basicConfig(
    filename='upload_debug.log',
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s'
)

@app.route("/")
def index():
    return render_template("index.html")


@app.route("/delete_file", methods=["POST"])
def delete_file():
    try:
        data = request.get_json()
        filename = data.get("filename")
        if not filename:
            return jsonify({"error": "缺少文件名"}), 400

        filepath = os.path.join("uploads", filename)
        if os.path.exists(filepath):
            os.remove(filepath)
            logging.info(f"✅ 已删除文件: {filename}")
            return jsonify({"success": True, "deleted": filename})
        else:
            return jsonify({"error": "文件不存在"}), 404

    except Exception as e:
        logging.exception("❌ 删除文件失败")
        return jsonify({"error": str(e)}), 500


@app.route("/uploaded_files", methods=["GET"])
def get_uploaded_files():
    """
    返回 uploads 文件夹中所有上传的 CSV/JSON 文件
    """
    try:
        files = list_uploaded_files()
        return jsonify({"files": files})
    except Exception as e:
        error_msg = f"获取上传文件列表失败: {str(e)}"
        logging.exception(error_msg)
        return jsonify({"error": error_msg}), 500



@app.route("/upload", methods=["POST"])
def upload():
    logging.info("收到上传请求")
    print("收到上传请求")

    if 'files' not in request.files:
        error_msg = "没有收到文件字段"
        logging.error(error_msg)
        return jsonify({"error": error_msg}), 400

    files = request.files.getlist("files")
    if not files:
        error_msg = "文件列表为空"
        logging.error(error_msg)
        return jsonify({"error": error_msg}), 400

    uploaded_filenames = []

    try:
        for file in files:
            if file and file.filename:
                logging.info(f"正在保存文件: {file.filename}")
                print(f"正在保存文件: {file.filename}")
                save_uploaded_file(file)
                uploaded_filenames.append(file.filename)

        logging.info(f"上传成功: {uploaded_filenames}")
        print(f"上传成功: {uploaded_filenames}")

        # ⚠️ 返回必须包含 files 字段！
        return jsonify({
            "success": True,
            "files": uploaded_filenames
        })

    except Exception as e:
        error_msg = f"保存文件时出错: {str(e)}"
        logging.exception(error_msg)
        return jsonify({"error": error_msg}), 500


@app.route("/files", methods=["GET"])
def get_files():
    files = list_uploaded_files()
    uploaded_filenames = [file.filename for file in files]
    return jsonify({"success": True, "files": uploaded_filenames})


def call_deepseek(messages):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}"
    }
    payload = {
        "model": "deepseek-chat",
        "messages": messages,
        "temperature": 0.2
    }
    response = requests.post(DEEPSEEK_API_URL, headers=headers, json=payload)
    response.raise_for_status()
    result = response.json()
    return result["choices"][0]["message"]["content"].strip()


def clean_chart_spec(raw_text):
    """更健壮的 JSON 清理函数"""
    # 尝试提取第一个完整的 JSON 对象
    try:
        # 查找可能的 JSON 起始位置
        start_idx = raw_text.find('{')
        if start_idx == -1:
            return raw_text.strip()

        # 尝试找到匹配的结束大括号
        brace_count = 0
        end_idx = -1
        for i in range(start_idx, len(raw_text)):
            if raw_text[i] == '{':
                brace_count += 1
            elif raw_text[i] == '}':
                brace_count -= 1
                if brace_count == 0:
                    end_idx = i
                    break

        # 如果找到匹配的结束大括号，提取 JSON
        if end_idx != -1:
            return raw_text[start_idx:end_idx + 1].strip()

        # 否则返回原始文本
        return raw_text.strip()
    except Exception:
        return raw_text.strip()

@app.route("/generate_from_selected", methods=["POST"])
def generate_from_selected():
    try:
        # 记录请求开始
        logging.info("="*50)
        logging.info("📌 开始处理图表生成请求")
        print("\n🔵 [图表生成] 收到新请求")

        # 获取并验证参数
        data = request.get_json()
        text = data.get("text", "").strip()
        selected_files = data.get("selected_files", [])
        
        logging.info(f"📝 用户描述: {text}")
        logging.info(f"📂 选择文件: {selected_files}")
        print(f"📝 用户描述: {text}")
        print(f"📂 选择文件: {selected_files}")

        if not text:
            error_msg = "❌ 错误: 未填写图表描述"
            logging.warning(error_msg)
            print(error_msg)
            return jsonify({"error": "请填写图表描述"}), 400

        if not selected_files:
            error_msg = "❌ 错误: 未选择文件"
            logging.warning(error_msg)
            print(error_msg)
            return jsonify({"error": "请选择至少一个文件"}), 400

        # 加载文件数据
        logging.info("⏳ 开始加载文件数据...")
        print("⏳ 加载文件数据...")
        dfs = load_selected_files(selected_files)
        logging.info(f"✅ 成功加载 {len(dfs)} 个文件")
        print(f"✅ 加载完成: 共 {len(dfs)} 个文件")

        # 准备数据提示
        combined = [df.to_dict(orient="records") for df in dfs]
        data_prompt = "\n\n".join([
            f"文件{i+1}（{fname}）数据：{json.dumps(data, ensure_ascii=False)}" 
            for i, (data, fname) in enumerate(zip(combined, selected_files))
        ])
        
        logging.info("📊 数据摘要:\n" + "\n".join([
            f"{fname}: {len(data)} 条记录" 
            for fname, data in zip(selected_files, combined)
        ]))
        print(f"📊 数据总记录数: {sum(len(d) for d in combined)}")

        # 调用AI生成
        messages = [
            {"role": "system", "content": "你是专业数据可视化助手"},
            {"role": "user", "content": f"""{data_prompt}
用户请求：{text}
⚠️ 注意事项：
- 横轴如包含数字，应按数值或时间顺序排序，而非按字符串顺序（例如 "1990-1"、"1990-2" 应排在 "1990-10" 前）；
- 日期字段如 "1990-1" 可视为年月，建议解析为标准格式如 "1990-01"；
- 只返回Vega-Lite JSON规范，不要包含任何Markdown代码块标记，避免使用未转义的特殊字符
"""}

        ]


        logging.info("🤖 正在调用DeepSeek API生成图表...")
        print("🔄 正在生成图表 (调用DeepSeek API)...")
        output = call_deepseek(messages)
        
        logging.info("🔧 清理API返回结果...")
        clean = clean_chart_spec(output)
        try:
            json_obj = json.loads(clean)
        except json.JSONDecodeError as e:
            # 记录错误详细信息
            logging.error(f"JSON解析错误: {str(e)}")
            logging.error(f"原始内容: {output[:500]}...")
            logging.error(f"清理后内容: {clean[:500]}...")
            raise ValueError("生成的图表规范无效，请调整描述重试")
        
        logging.info(f"🎉 图表生成成功！规范大小: {len(clean)} 字符")
        print(f"✅ 图表生成完成！规范大小: {len(clean)} 字符")
        
        # 记录结束
        logging.info("="*50 + "\n")
        print("="*50 + "\n")

        return jsonify({
            "spec": json.dumps(json_obj, ensure_ascii=False, indent=2),
            "debug": {
                "input_files": selected_files,
                "data_records": sum(len(d) for d in combined),
                "spec_size": len(clean)
            }
        })

    except Exception as e:
        error_msg = f"❌ 生成过程中出错: {str(e)}"
        logging.exception(error_msg)
        print(error_msg)
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(debug=True)